Skip to content

Fit PCA in batches to overcome LAPACK memory overflow - #4869

Open
kushaangupta wants to merge 2 commits into
SpikeInterface:mainfrom
kushaangupta:fix/global-pca-fit-batches
Open

Fit PCA in batches to overcome LAPACK memory overflow#4869
kushaangupta wants to merge 2 commits into
SpikeInterface:mainfrom
kushaangupta:fix/global-pca-fit-batches

Conversation

@kushaangupta

@kushaangupta kushaangupta commented Oct 14, 2026

Copy link
Copy Markdown
Contributor

Fixes #4868

  • Add max_fit_batch_memory="64Mi" for PCA fitting. The 64 MiB default is a proposed input-buffer budget, not a measured optimum or a bound on total process RAM.
  • Read bounded groups of waveforms through slice_rows() before reshaping. This supports NumPy and Zarr storage and avoids first copying the whole unit.
  • Include sparse padding when budgeting reads and exclude padded channels from the PCA rows.
  • Also cap the augmented SVD matrix below the 32-bit LAPACK indexing limit, accounting for the previous components and mean-correction row.
  • Preserve input order, current small-unit update boundaries, whitening, component count and the existing too-few-rows rule. Balance batch sizes to retain a short remainder; reject a budget that cannot support valid updates.

@kushaangupta
Bound global PCA fitting batches to avoid LAPACK index overflow
4ca5171 
Add a configurable 64 MiB waveform input budget for by_channel_global fitting, read batches before reshaping, and cap augmented SVD dimensions below the 32-bit LAPACK limit. Preserve input ordering, sparse channel handling and existing update boundaries for small units. Document numerical implications and cover NumPy/Zarr reads, complete input coverage and extension round trips.
@kushaangupta
Bound local and concatenated PCA updates as well as global fitting
f3481cc 
Share waveform batching with mode-specific PCA row and feature counts. Preserve each channel model update order and reuse local workers across bounded batches. Cover dense and sparse NumPy/Zarr fitting across all modes, and compare serial with spawn-worker local results.
@alejoe91 alejoe91 added the postprocessing Related to postprocessing module label Oct 15, 2026
@JoeZiminski

JoeZiminski commented Oct 21, 2026

Copy link
Copy Markdown
Contributor

Thanks for this @kushaangupta, good catch that's funny that we are hitting the 32-bit indexing limit on these matrices 😂

This is a nice approach but I wonder if it could be more targetted to fix only the bug at hand, rather than extending to provide general memory handling. Fundamentally the main thing we want to avoid is passing too large matrices to partial fit and batch them instead, e.g. I think this could be achieved with something like the below helper:

import numpy as np
from sklearn.utils import gen_batches

def safe_partial_fit(pca, waveforms):
    n_samples, n_features = waveforms.shape

    # LAPACK indexing is limited to 2**31 - 1 elements. Allow for the extra
    # rows IncrementalPCA adds internally before computing the SVD.
    max_rows = np.iinfo(np.int32).max // n_features - (pca.n_components + 1)

    for batch in gen_batches(
        n_samples,
        max_rows,
        min_batch_size=pca.n_components,
    ):
        pca.partial_fit(waveforms[batch])

This is basically the approach IncrementalPCA is taking.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

postprocessing Related to postprocessing module

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

PCA fitting can overflow LAPACK's 32-bit indexing when fitting all spikes

3 participants

@kushaangupta @JoeZiminski @alejoe91